{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2e0f5ca-4893-4cfa-b6d4-0447f0bc314b",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Preliminary BERT results, 19 April 2023\n",
    "\n",
    "This notebook contains the experiment that was presented at the 19 April 2023 participant sync. The first few cells present some basic data transformations to prepare the data for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "688b057f-17b0-41bd-bca1-2fbfd2f4cfc9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>order</th>\n",
       "      <th>disposition</th>\n",
       "      <th>mappings</th>\n",
       "      <th>annotator</th>\n",
       "      <th>reference</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>has overwritten the function pointer in the ex...</td>\n",
       "      <td></td>\n",
       "      <td>accept</td>\n",
       "      <td>[{'technique_name': 'Extra Window Memory Injec...</td>\n",
       "      <td>[{'organization_name': 'MITRE', 'annotator_nam...</td>\n",
       "      <td>[[{'url': 'https://recon.cx/2018/brussels/reso...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>overwrites Explorers Shell_TrayWnd extra windo...</td>\n",
       "      <td></td>\n",
       "      <td>accept</td>\n",
       "      <td>[{'technique_name': 'Extra Window Memory Injec...</td>\n",
       "      <td>[{'organization_name': 'MITRE', 'annotator_nam...</td>\n",
       "      <td>[[{'url': 'https://www.malwaretech.com/2013/08...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td></td>\n",
       "      <td>accept</td>\n",
       "      <td>[{'technique_name': 'Scheduled Task', 'attack_...</td>\n",
       "      <td>[{'organization_name': 'MITRE', 'annotator_nam...</td>\n",
       "      <td>[[{'url': 'https://www.microsoft.com/security/...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>has the ability to launch scheduled tasks to e...</td>\n",
       "      <td></td>\n",
       "      <td>accept</td>\n",
       "      <td>[{'technique_name': 'Scheduled Task', 'attack_...</td>\n",
       "      <td>[{'organization_name': 'MITRE', 'annotator_nam...</td>\n",
       "      <td>[[{'url': 'https://www.crowdstrike.com/blog/ca...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td></td>\n",
       "      <td>accept</td>\n",
       "      <td>[{'technique_name': 'Scheduled Task', 'attack_...</td>\n",
       "      <td>[{'organization_name': 'MITRE', 'annotator_nam...</td>\n",
       "      <td>[[{'url': 'https://go.crowdstrike.com/rs/281-O...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24599</th>\n",
       "      <td>\"My God\" was one of the first songs recorded b...</td>\n",
       "      <td>12583</td>\n",
       "      <td>accept</td>\n",
       "      <td>[]</td>\n",
       "      <td>[{'organization_name': 'unknown', 'annotator_n...</td>\n",
       "      <td>[{'url': 'https://github.com/center-for-threat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24600</th>\n",
       "      <td>It initially had seven students.</td>\n",
       "      <td>12584</td>\n",
       "      <td>accept</td>\n",
       "      <td>[]</td>\n",
       "      <td>[{'organization_name': 'unknown', 'annotator_n...</td>\n",
       "      <td>[{'url': 'https://github.com/center-for-threat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24601</th>\n",
       "      <td>Vellarikundu is a hillside town and taluk head...</td>\n",
       "      <td>12585</td>\n",
       "      <td>accept</td>\n",
       "      <td>[]</td>\n",
       "      <td>[{'organization_name': 'unknown', 'annotator_n...</td>\n",
       "      <td>[{'url': 'https://github.com/center-for-threat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24602</th>\n",
       "      <td>This earned the score a parental advisory warn...</td>\n",
       "      <td>12586</td>\n",
       "      <td>accept</td>\n",
       "      <td>[]</td>\n",
       "      <td>[{'organization_name': 'unknown', 'annotator_n...</td>\n",
       "      <td>[{'url': 'https://github.com/center-for-threat...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24603</th>\n",
       "      <td>Acanthocardamum is a monotypic genus in the fa...</td>\n",
       "      <td>12587</td>\n",
       "      <td>accept</td>\n",
       "      <td>[]</td>\n",
       "      <td>[{'organization_name': 'unknown', 'annotator_n...</td>\n",
       "      <td>[{'url': 'https://github.com/center-for-threat...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>24604 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    text  order disposition   \n",
       "0      has overwritten the function pointer in the ex...             accept  \\\n",
       "1      overwrites Explorers Shell_TrayWnd extra windo...             accept   \n",
       "2      has used scheduled tasks to maintain persistence.             accept   \n",
       "3      has the ability to launch scheduled tasks to e...             accept   \n",
       "4      has used scheduled tasks to maintain persistence.             accept   \n",
       "...                                                  ...    ...         ...   \n",
       "24599  \"My God\" was one of the first songs recorded b...  12583      accept   \n",
       "24600                   It initially had seven students.  12584      accept   \n",
       "24601  Vellarikundu is a hillside town and taluk head...  12585      accept   \n",
       "24602  This earned the score a parental advisory warn...  12586      accept   \n",
       "24603  Acanthocardamum is a monotypic genus in the fa...  12587      accept   \n",
       "\n",
       "                                                mappings   \n",
       "0      [{'technique_name': 'Extra Window Memory Injec...  \\\n",
       "1      [{'technique_name': 'Extra Window Memory Injec...   \n",
       "2      [{'technique_name': 'Scheduled Task', 'attack_...   \n",
       "3      [{'technique_name': 'Scheduled Task', 'attack_...   \n",
       "4      [{'technique_name': 'Scheduled Task', 'attack_...   \n",
       "...                                                  ...   \n",
       "24599                                                 []   \n",
       "24600                                                 []   \n",
       "24601                                                 []   \n",
       "24602                                                 []   \n",
       "24603                                                 []   \n",
       "\n",
       "                                               annotator   \n",
       "0      [{'organization_name': 'MITRE', 'annotator_nam...  \\\n",
       "1      [{'organization_name': 'MITRE', 'annotator_nam...   \n",
       "2      [{'organization_name': 'MITRE', 'annotator_nam...   \n",
       "3      [{'organization_name': 'MITRE', 'annotator_nam...   \n",
       "4      [{'organization_name': 'MITRE', 'annotator_nam...   \n",
       "...                                                  ...   \n",
       "24599  [{'organization_name': 'unknown', 'annotator_n...   \n",
       "24600  [{'organization_name': 'unknown', 'annotator_n...   \n",
       "24601  [{'organization_name': 'unknown', 'annotator_n...   \n",
       "24602  [{'organization_name': 'unknown', 'annotator_n...   \n",
       "24603  [{'organization_name': 'unknown', 'annotator_n...   \n",
       "\n",
       "                                               reference  \n",
       "0      [[{'url': 'https://recon.cx/2018/brussels/reso...  \n",
       "1      [[{'url': 'https://www.malwaretech.com/2013/08...  \n",
       "2      [[{'url': 'https://www.microsoft.com/security/...  \n",
       "3      [[{'url': 'https://www.crowdstrike.com/blog/ca...  \n",
       "4      [[{'url': 'https://go.crowdstrike.com/rs/281-O...  \n",
       "...                                                  ...  \n",
       "24599  [{'url': 'https://github.com/center-for-threat...  \n",
       "24600  [{'url': 'https://github.com/center-for-threat...  \n",
       "24601  [{'url': 'https://github.com/center-for-threat...  \n",
       "24602  [{'url': 'https://github.com/center-for-threat...  \n",
       "24603  [{'url': 'https://github.com/center-for-threat...  \n",
       "\n",
       "[24604 rows x 6 columns]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "data_path = '/projects/TRAM2023/tram-private/data/training/refreshed_dataset_march_2023.json'\n",
    "with open(data_path) as f:\n",
    "    data = json.loads(f.read())\n",
    "    \n",
    "import pandas as pd\n",
    "raw = pd.DataFrame(data['sentences'])\n",
    "raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cff92e03-f726-4206-bc38-ecc80fb09361",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>technique_name</th>\n",
       "      <th>attack_id</th>\n",
       "      <th>confidence</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Extra Window Memory Injection</td>\n",
       "      <td>T1055.011</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Extra Window Memory Injection</td>\n",
       "      <td>T1055.011</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Scheduled Task</td>\n",
       "      <td>T1053.005</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Scheduled Task</td>\n",
       "      <td>T1053.005</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Scheduled Task</td>\n",
       "      <td>T1053.005</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13536</th>\n",
       "      <td>Emond</td>\n",
       "      <td>T1546.014</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13537</th>\n",
       "      <td>Control Panel</td>\n",
       "      <td>T1218.002</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13538</th>\n",
       "      <td>Control Panel</td>\n",
       "      <td>T1218.002</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13539</th>\n",
       "      <td>Application Shimming</td>\n",
       "      <td>T1546.011</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13540</th>\n",
       "      <td>Application Shimming</td>\n",
       "      <td>T1546.011</td>\n",
       "      <td>100.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>13541 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                      technique_name  attack_id confidence\n",
       "0      Extra Window Memory Injection  T1055.011      100.0\n",
       "1      Extra Window Memory Injection  T1055.011      100.0\n",
       "2                     Scheduled Task  T1053.005      100.0\n",
       "3                     Scheduled Task  T1053.005      100.0\n",
       "4                     Scheduled Task  T1053.005      100.0\n",
       "...                              ...        ...        ...\n",
       "13536                          Emond  T1546.014      100.0\n",
       "13537                  Control Panel  T1218.002      100.0\n",
       "13538                  Control Panel  T1218.002      100.0\n",
       "13539           Application Shimming  T1546.011      100.0\n",
       "13540           Application Shimming  T1546.011      100.0\n",
       "\n",
       "[13541 rows x 3 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mappings = raw['mappings'].explode().dropna().apply(pd.Series)\n",
    "mappings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d7a60415-c81c-4fd5-a240-949a435b3f1b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>attack_id</th>\n",
       "      <th>subclass_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>has overwritten the function pointer in the ex...</td>\n",
       "      <td>T1055</td>\n",
       "      <td>011</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>overwrites Explorers Shell_TrayWnd extra windo...</td>\n",
       "      <td>T1055</td>\n",
       "      <td>011</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>has the ability to launch scheduled tasks to e...</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24599</th>\n",
       "      <td>\"My God\" was one of the first songs recorded b...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24600</th>\n",
       "      <td>It initially had seven students.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24601</th>\n",
       "      <td>Vellarikundu is a hillside town and taluk head...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24602</th>\n",
       "      <td>This earned the score a parental advisory warn...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24603</th>\n",
       "      <td>Acanthocardamum is a monotypic genus in the fa...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>24604 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    text attack_id subclass_id\n",
       "0      has overwritten the function pointer in the ex...     T1055         011\n",
       "1      overwrites Explorers Shell_TrayWnd extra windo...     T1055         011\n",
       "2      has used scheduled tasks to maintain persistence.     T1053         005\n",
       "3      has the ability to launch scheduled tasks to e...     T1053         005\n",
       "4      has used scheduled tasks to maintain persistence.     T1053         005\n",
       "...                                                  ...       ...         ...\n",
       "24599  \"My God\" was one of the first songs recorded b...       NaN         NaN\n",
       "24600                   It initially had seven students.       NaN         NaN\n",
       "24601  Vellarikundu is a hillside town and taluk head...       NaN         NaN\n",
       "24602  This earned the score a parental advisory warn...       NaN         NaN\n",
       "24603  Acanthocardamum is a monotypic genus in the fa...       NaN         NaN\n",
       "\n",
       "[24604 rows x 3 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.concat((raw['text'], mappings['attack_id'].str.extract(r\"(?P<attack_id>T\\d+)\\.(?P<subclass_id>\\d+)\")), axis=1)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9bf4b70d-d13a-4293-8534-10560cc222b1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "attack_id\n",
       "NaN      17932\n",
       "T1059      706\n",
       "T1071      397\n",
       "T1070      368\n",
       "T1547      337\n",
       "         ...  \n",
       "T1011        1\n",
       "T1499        1\n",
       "T1216        1\n",
       "T1597        1\n",
       "T1601        1\n",
       "Name: count, Length: 89, dtype: int64"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['attack_id'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "585d96e2-ad7a-423f-905a-3188f4fc7522",
   "metadata": {},
   "source": [
    "Here we see that unlabeled text segments are significantly more frequent than any individual technique. We will include 1000 of them in the data for this experiment, along with all instances of these ATT&CK techniques.\n",
    "\n",
    "'T1041', 'T1106', 'T1082', 'T1033', 'T1112', 'T1070', 'T1090', 'T1021', 'T1218', 'T1095', 'T1548', 'T1053', 'T1071', 'T1574', 'T1562', 'T1204', 'T1012', 'T1140', 'T1055', 'T1105', 'T1552', 'T1486', 'T1083', 'T1078', 'T1047', 'T1190', 'T1543', 'T1113', 'T1003', 'T1059', 'T1057', 'T1027', 'T1219', 'T1036', 'T1005'\n",
    "\n",
    "Note that not all of these techniques are present in the available data, but for each that is present, we use every instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a68d1eb7-f554-4d57-bbd8-4f17b8b50383",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>attack_id</th>\n",
       "      <th>subclass_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>has overwritten the function pointer in the ex...</td>\n",
       "      <td>T1055</td>\n",
       "      <td>011</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>overwrites Explorers Shell_TrayWnd extra windo...</td>\n",
       "      <td>T1055</td>\n",
       "      <td>011</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>has the ability to launch scheduled tasks to e...</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>has used scheduled tasks to maintain persistence.</td>\n",
       "      <td>T1053</td>\n",
       "      <td>005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23524</th>\n",
       "      <td>In 2017, Waterhouse shifted to Australian Rule...</td>\n",
       "      <td>none</td>\n",
       "      <td>none</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21207</th>\n",
       "      <td>The district was formed in 1973 under the Loca...</td>\n",
       "      <td>none</td>\n",
       "      <td>none</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6738</th>\n",
       "      <td>has a command to list the victim's processes.</td>\n",
       "      <td>none</td>\n",
       "      <td>none</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19345</th>\n",
       "      <td>Neophile or Neophiliac, a term popularized by ...</td>\n",
       "      <td>none</td>\n",
       "      <td>none</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14710</th>\n",
       "      <td>The company sells over 400 products, including...</td>\n",
       "      <td>none</td>\n",
       "      <td>none</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>4527 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    text attack_id subclass_id\n",
       "0      has overwritten the function pointer in the ex...     T1055         011\n",
       "1      overwrites Explorers Shell_TrayWnd extra windo...     T1055         011\n",
       "2      has used scheduled tasks to maintain persistence.     T1053         005\n",
       "3      has the ability to launch scheduled tasks to e...     T1053         005\n",
       "4      has used scheduled tasks to maintain persistence.     T1053         005\n",
       "...                                                  ...       ...         ...\n",
       "23524  In 2017, Waterhouse shifted to Australian Rule...      none        none\n",
       "21207  The district was formed in 1973 under the Loca...      none        none\n",
       "6738       has a command to list the victim's processes.      none        none\n",
       "19345  Neophile or Neophiliac, a term popularized by ...      none        none\n",
       "14710  The company sells over 400 products, including...      none        none\n",
       "\n",
       "[4527 rows x 3 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classes_of_interest = ['T1041', 'T1106', 'T1082', 'T1033', 'T1112', 'T1070', 'T1090', 'T1021', 'T1218', 'T1095', 'T1548', 'T1053', 'T1071', 'T1574', 'T1562', 'T1204', 'T1012', 'T1140', 'T1055', 'T1105', 'T1552', 'T1486', 'T1083', 'T1078', 'T1047', 'T1190', 'T1543', 'T1113', 'T1003', 'T1059', 'T1057', 'T1027', 'T1219', 'T1036', 'T1005']\n",
    "positive_data = df[df['attack_id'].isin(classes_of_interest)]\n",
    "negative_data = df[df['attack_id'].isna()].sample(1000).fillna('none')\n",
    "data = pd.concat((positive_data, negative_data))\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4df46da6-7001-478f-a1ba-659f92460a95",
   "metadata": {},
   "source": [
    "For this preliminary experiment, we will use the model `bert-base-cased`, which is the BERT model provided as a reference in the original BERT paper by Devlin et al."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "29b205cf-32ba-4ea7-903d-348d8e5c9e3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import transformers\n",
    "import torch\n",
    "\n",
    "cuda = torch.device('cuda')\n",
    "tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-cased', max_length=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a0914d75-240a-4087-9f61-273fc734294f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  101,  1144,  1166,  ...,     0,     0,     0],\n",
       "        [  101,  1166,  2246,  ...,     0,     0,     0],\n",
       "        [  101,  1144,  1215,  ...,     0,     0,     0],\n",
       "        ...,\n",
       "        [  101,  1144,   170,  ...,     0,     0,     0],\n",
       "        [  101, 14521, 27008,  ...,     0,     0,     0],\n",
       "        [  101,  1109,  1419,  ...,     0,     0,     0]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_tokens = tokenizer(data['text'].tolist(), return_tensors='pt', padding='max_length', truncation=True, max_length=512).input_ids\n",
    "x_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "33ee7ad2-b923-4fa3-99cf-0e99d6aa3bf1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4527, 512])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_tokens.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b860c004-15b7-48d6-a00f-43e0296c139c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'T1055',\n",
       " 1: 'T1053',\n",
       " 2: 'T1021',\n",
       " 3: 'T1218',\n",
       " 4: 'T1027',\n",
       " 5: 'T1574',\n",
       " 6: 'T1059',\n",
       " 7: 'T1036',\n",
       " 8: 'T1548',\n",
       " 9: 'T1003',\n",
       " 10: 'T1071',\n",
       " 11: 'T1552',\n",
       " 12: 'T1204',\n",
       " 13: 'T1562',\n",
       " 14: 'T1543',\n",
       " 15: 'T1070',\n",
       " 16: 'T1078',\n",
       " 17: 'T1090',\n",
       " 18: 'none'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_to_label = dict(enumerate(data['attack_id'].unique()))\n",
    "index_to_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f539c6ec-3cb8-478a-b541-51902a83ae0b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'T1055': 0,\n",
       " 'T1053': 1,\n",
       " 'T1021': 2,\n",
       " 'T1218': 3,\n",
       " 'T1027': 4,\n",
       " 'T1574': 5,\n",
       " 'T1059': 6,\n",
       " 'T1036': 7,\n",
       " 'T1548': 8,\n",
       " 'T1003': 9,\n",
       " 'T1071': 10,\n",
       " 'T1552': 11,\n",
       " 'T1204': 12,\n",
       " 'T1562': 13,\n",
       " 'T1543': 14,\n",
       " 'T1070': 15,\n",
       " 'T1078': 16,\n",
       " 'T1090': 17,\n",
       " 'none': 18}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_to_index = {label: index for index, label in index_to_label.items()}\n",
    "label_to_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "579bce5e-9070-4715-afc0-56db487a7b44",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0,  0,  1,  ..., 18, 18, 18])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_all = torch.Tensor(data['attack_id'].replace(label_to_index).to_numpy()).to(int)\n",
    "y_all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c9cc5d8-c181-4f21-a994-c23ea01f55be",
   "metadata": {},
   "source": [
    "We split the data 80/20 between train and test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ccb349da-56a9-40a7-8efc-9b61ff212d88",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "x_train, x_test, y_train, y_test = train_test_split(x_tokens, y_all, test_size=0.2, shuffle=True)\n",
    "\n",
    "def _load_data(x, y, batch_size=10):\n",
    "    x_len, y_len = x.shape[0], y.shape[0]\n",
    "    assert x_len == y_len\n",
    "    for i in range(0, x_len, batch_size):\n",
    "        slc = slice(i, i + batch_size)\n",
    "        yield x[slc].to(cuda), y[slc].to(cuda)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eefa82eb-e5b9-4e98-9d74-ae41699af35b",
   "metadata": {},
   "source": [
    "Each vector representing an instance is padded with trailing zeros, so that every vector is the same length, 512. In the expression `bert(x, attention_mask=x.ne(0).to(int), labels=y)`, the attention mask indicates that all non-zero elements are important. (`ne` means \"not equal\", so `x.ne(0).to(int)` returns a binary array of 1s for non-zero elements of x, else 0.)\n",
    "\n",
    "We use a learning rate of $2 \\times 10^{-5}$, $\\epsilon = 1 \\times 10^{-8}$, and a batch size of 10, for five epochs.\n",
    "\n",
    "Note that the loss function is (rather unusually) called internally to `bert`. Looking at the source code for `transformers.BertForSequenceClassification`, the loss function that is used in this experiment is cross entropy loss. (My exploration of the source code revealed that `BertForSequenceClassification` will also support multi-label classification, which is good news as we expect to do that later in the project.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d343ac68-0a1f-4160-82ae-b1a5a9fbd0d0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertForSequenceClassification: ['cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "363it [01:52,  3.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 0 loss: 1.6576010820868916\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "363it [01:54,  3.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1 loss: 0.5818408697474101\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "363it [01:57,  3.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 2 loss: 0.2602325600902896\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "363it [01:57,  3.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 3 loss: 0.12686511959377728\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "363it [01:57,  3.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 4 loss: 0.07342217404128516\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from statistics import mean\n",
    "\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import AdamW\n",
    "\n",
    "bert = transformers.BertForSequenceClassification.from_pretrained(\n",
    "    'bert-base-cased',\n",
    "    num_labels=data['attack_id'].nunique(),\n",
    "    output_attentions=False, \n",
    "    output_hidden_states=False,\n",
    ")\n",
    "bert.train().to(cuda)\n",
    "optim = AdamW(bert.parameters(), lr=2e-5, eps=1e-8)\n",
    "\n",
    "for epoch in range(5):\n",
    "    epoch_losses = []\n",
    "    for x, y in tqdm(_load_data(x_train, y_train, batch_size=10)):\n",
    "        bert.zero_grad()\n",
    "        out = bert(x, attention_mask=x.ne(0).to(int), labels=y)\n",
    "        epoch_losses.append(out.loss.item())\n",
    "        out.loss.backward()\n",
    "        optim.step()\n",
    "    print(f\"epoch {epoch} loss: {mean(epoch_losses)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2641161d-7e3c-4e82-946c-acc2c784588f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import precision_recall_fscore_support as calculate_score\n",
    "\n",
    "bert.eval()\n",
    "\n",
    "batch_size = 20\n",
    "preds = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in range(0, x_test.shape[0], batch_size):\n",
    "        x = x_test[i : i + batch_size].to(cuda)\n",
    "        out = bert(x, attention_mask=x.ne(0).to(int))\n",
    "        preds.extend(out.logits.argmax(-1).to('cpu').numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6de66fa-a4c4-4e0f-b0a5-c41548061cb7",
   "metadata": {},
   "source": [
    "As we are doing single-label classification, the prediction is considered to be the argmax of the logits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "35f5ab11-d79a-4912-9758-b6615f9c405f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>P</th>\n",
       "      <th>R</th>\n",
       "      <th>F1</th>\n",
       "      <th>#</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>T1003</th>\n",
       "      <td>0.906250</td>\n",
       "      <td>0.828571</td>\n",
       "      <td>0.865672</td>\n",
       "      <td>35.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1021</th>\n",
       "      <td>0.913043</td>\n",
       "      <td>0.750000</td>\n",
       "      <td>0.823529</td>\n",
       "      <td>28.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1027</th>\n",
       "      <td>0.857143</td>\n",
       "      <td>0.789474</td>\n",
       "      <td>0.821918</td>\n",
       "      <td>38.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1036</th>\n",
       "      <td>0.925000</td>\n",
       "      <td>0.925000</td>\n",
       "      <td>0.925000</td>\n",
       "      <td>40.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1053</th>\n",
       "      <td>0.939394</td>\n",
       "      <td>0.968750</td>\n",
       "      <td>0.953846</td>\n",
       "      <td>32.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1055</th>\n",
       "      <td>0.888889</td>\n",
       "      <td>0.888889</td>\n",
       "      <td>0.888889</td>\n",
       "      <td>27.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1059</th>\n",
       "      <td>0.917647</td>\n",
       "      <td>0.981132</td>\n",
       "      <td>0.948328</td>\n",
       "      <td>159.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1070</th>\n",
       "      <td>0.962500</td>\n",
       "      <td>0.962500</td>\n",
       "      <td>0.962500</td>\n",
       "      <td>80.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1071</th>\n",
       "      <td>0.945946</td>\n",
       "      <td>0.958904</td>\n",
       "      <td>0.952381</td>\n",
       "      <td>73.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1078</th>\n",
       "      <td>0.700000</td>\n",
       "      <td>0.875000</td>\n",
       "      <td>0.777778</td>\n",
       "      <td>8.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1090</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.636364</td>\n",
       "      <td>0.777778</td>\n",
       "      <td>11.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1204</th>\n",
       "      <td>0.900000</td>\n",
       "      <td>0.947368</td>\n",
       "      <td>0.923077</td>\n",
       "      <td>38.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1218</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.829268</td>\n",
       "      <td>0.906667</td>\n",
       "      <td>41.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1543</th>\n",
       "      <td>0.931034</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.964286</td>\n",
       "      <td>27.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1548</th>\n",
       "      <td>0.785714</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.880000</td>\n",
       "      <td>11.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1552</th>\n",
       "      <td>0.818182</td>\n",
       "      <td>0.750000</td>\n",
       "      <td>0.782609</td>\n",
       "      <td>12.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1562</th>\n",
       "      <td>0.888889</td>\n",
       "      <td>0.695652</td>\n",
       "      <td>0.780488</td>\n",
       "      <td>23.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>T1574</th>\n",
       "      <td>0.904762</td>\n",
       "      <td>0.791667</td>\n",
       "      <td>0.844444</td>\n",
       "      <td>24.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>none</th>\n",
       "      <td>0.865385</td>\n",
       "      <td>0.904523</td>\n",
       "      <td>0.884521</td>\n",
       "      <td>199.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>(micro)</th>\n",
       "      <td>0.906181</td>\n",
       "      <td>0.906181</td>\n",
       "      <td>0.906181</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>(macro)</th>\n",
       "      <td>0.897357</td>\n",
       "      <td>0.867530</td>\n",
       "      <td>0.877037</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                P         R        F1      #\n",
       "T1003    0.906250  0.828571  0.865672   35.0\n",
       "T1021    0.913043  0.750000  0.823529   28.0\n",
       "T1027    0.857143  0.789474  0.821918   38.0\n",
       "T1036    0.925000  0.925000  0.925000   40.0\n",
       "T1053    0.939394  0.968750  0.953846   32.0\n",
       "T1055    0.888889  0.888889  0.888889   27.0\n",
       "T1059    0.917647  0.981132  0.948328  159.0\n",
       "T1070    0.962500  0.962500  0.962500   80.0\n",
       "T1071    0.945946  0.958904  0.952381   73.0\n",
       "T1078    0.700000  0.875000  0.777778    8.0\n",
       "T1090    1.000000  0.636364  0.777778   11.0\n",
       "T1204    0.900000  0.947368  0.923077   38.0\n",
       "T1218    1.000000  0.829268  0.906667   41.0\n",
       "T1543    0.931034  1.000000  0.964286   27.0\n",
       "T1548    0.785714  1.000000  0.880000   11.0\n",
       "T1552    0.818182  0.750000  0.782609   12.0\n",
       "T1562    0.888889  0.695652  0.780488   23.0\n",
       "T1574    0.904762  0.791667  0.844444   24.0\n",
       "none     0.865385  0.904523  0.884521  199.0\n",
       "(micro)  0.906181  0.906181  0.906181    NaN\n",
       "(macro)  0.897357  0.867530  0.877037    NaN"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def calculate_scores_df(actual: list[str], predicted: list[str]):\n",
    "    scores = calculate_score(actual, predicted)\n",
    "    scores_df = pd.DataFrame(scores).T\n",
    "    scores_df.columns = ['P', 'R', 'F1', '#']\n",
    "    scores_df.index = sorted(set(actual) | set(predicted))\n",
    "    scores_df.loc['(micro)'] = calculate_score(actual, predicted, average='micro')\n",
    "    scores_df.loc['(macro)'] = calculate_score(actual, predicted, average='macro')\n",
    "    return scores_df\n",
    "\n",
    "y_test_list = pd.Series(y_test.tolist()).replace(index_to_label)\n",
    "preds_list = pd.Series(preds).replace(index_to_label)\n",
    "\n",
    "results = calculate_scores_df(y_test_list, preds_list)\n",
    "results"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
